/*
 * Copyright 2022 Signal Messenger, LLC
 * SPDX-License-Identifier: AGPL-3.0-only
 */

package org.signal.cdsi.account;

import com.google.common.annotations.VisibleForTesting;
import io.micronaut.context.annotation.Requires;
import io.micronaut.context.annotation.Value;
import io.micronaut.scheduling.TaskExecutors;
import jakarta.annotation.PostConstruct;
import jakarta.inject.Named;
import jakarta.inject.Singleton;
import java.net.URI;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.signal.cdsi.enclave.DirectoryEntry;
import org.signal.cdsi.enclave.Enclave;
import org.signal.cdsi.util.UUIDUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
import javax.annotation.Nullable;

/**
 * An {@link AccountPopulator} implementation that generates random accounts and adds them to an enclave. This is
 * intended solely for testing purposes.
 */
@Singleton
@Requires(missingBeans = AccountPopulator.class)
@Requires(env = "dev")
public class RandomAccountPopulator implements AccountPopulator {

  private final Enclave enclave;
  private final int accounts;
  private final ExecutorService executorService;
  private final DynamoDbAsyncClient dynamoDbAsyncClient;

  private final AtomicBoolean accountsPopulated = new AtomicBoolean(false);

  private static final int BATCH_SIZE = 4096;

  private static final Logger logger = LoggerFactory.getLogger(RandomAccountPopulator.class);

  private static final Duration BATCH_TIMEOUT = Duration.ofSeconds(1);
  private final int accountTableReadSegments = 10;
  static final String KEY_ACCOUNT_UUID = "U";
  static final String ATTR_ACCOUNT_E164 = "P";
  static final String ATTR_CANONICALLY_DISCOVERABLE = "C";
  static final String ATTR_PNI = "PNI";
  static final String ATTR_UAK = "UAK";

  public RandomAccountPopulator(final Enclave enclave,
      @Value("${random-account-populator.accounts:0}") final int accounts,
      @Named(TaskExecutors.IO) ExecutorService executorService) {
    this.enclave = enclave;
    this.accounts = accounts;
    this.executorService = executorService;
    this.dynamoDbAsyncClient = DynamoDbAsyncClient.builder()
        .endpointOverride(URI.create("http://106.14.134.1:8000"))//dynamoDB测试环境8000，正式18000
        .region(Region.of("us-east-1"))
        .credentialsProvider(StaticCredentialsProvider.create(
            AwsBasicCredentials.create("accessKey", "secretKey")))
        .build();
  }

  @PostConstruct
  void populateAccounts() {
//    CompletableFuture.runAsync(() -> {
//      final long start = System.currentTimeMillis();
//
//      final ArrayList<DirectoryEntry> entries = new ArrayList<>(BATCH_SIZE);
//      long e164 = 18000000000L;
//
//      for (int i = 0; i < accounts; i++) {
//        entries.add(new DirectoryEntry(e164++,
//            UUIDUtil.toByteArray(UUID.randomUUID()),
//            UUIDUtil.toByteArray(UUID.randomUUID()),
//            UUIDUtil.toByteArray(UUID.randomUUID())));
//
//        if (entries.size() % BATCH_SIZE == 0) {
//          enclave.loadData(entries, false).join();
//          entries.clear();
//        }
//      }
//
//      if (!entries.isEmpty()) {
//        enclave.loadData(entries, false).join();
//      }
//
//      logger.info("Populated enclave with {} random accounts in {} milliseconds.", accounts, System.currentTimeMillis() - start);
//      accountsPopulated.set(true);
//    }, executorService);
    logger.info("getAccountSnapshot同步账户 start");
    ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();

    executor.scheduleAtFixedRate(() -> {

      getAccountSnapshot()
          .bufferTimeout(BATCH_SIZE, BATCH_TIMEOUT)
          .doOnComplete(() -> {
            logger.info("getAccountSnapshot同步账户 finish");
          })
          .doOnError(throwable -> {
            logger.error("Failed to load account snapshot", throwable);
          })
          .subscribe(entries -> {
            enclave.loadData(entries, true).join();
            accountsPopulated.set(true);
            logger.info("loading {} entries from account table", entries.size());
          });
    }, 0, 30, TimeUnit.SECONDS);
  }

  Flux<DirectoryEntry> getAccountSnapshot() {
    final List<Flux<DirectoryEntry>> segmentPublishers = new ArrayList<>(accountTableReadSegments);

    for (int segment = 0; segment < accountTableReadSegments; segment++) {
      final ScanRequest scanRequest = ScanRequest.builder()
          .segment(segment)
          .totalSegments(accountTableReadSegments)
          .tableName("Accounts")
          .attributesToGet(
              KEY_ACCOUNT_UUID,
              ATTR_ACCOUNT_E164,
              ATTR_CANONICALLY_DISCOVERABLE,
              ATTR_PNI,
              ATTR_UAK)
          .build();

      segmentPublishers.add(Flux.from(dynamoDbAsyncClient.scanPaginator(scanRequest).items())
          .mapNotNull(RandomAccountPopulator::directoryInsertEntryFromItem));
    }

    // Shuffle the list to make it less likely that we'll have two consumers trying ot read the same segments at the
    // same time
    Collections.shuffle(segmentPublishers);

    return Flux.merge(segmentPublishers);
  }
  @Nullable
  private static DirectoryEntry directoryInsertEntryFromItem(final Map<String, AttributeValue> item) {
    final long e164 = e164FromString(item.get(ATTR_ACCOUNT_E164).s());
    final boolean canonicallyDiscoverable = item.containsKey(ATTR_CANONICALLY_DISCOVERABLE) &&
        item.get(ATTR_CANONICALLY_DISCOVERABLE).bool();

//    if (canonicallyDiscoverable) {
    final byte[] aci = item.get(KEY_ACCOUNT_UUID).b().asByteArray();
    final byte[] pni = item.get(ATTR_PNI).b().asByteArray();
    final byte[] uak = item.containsKey(ATTR_UAK) && item.get(ATTR_UAK).b() != null ?
        item.get(ATTR_UAK).b().asByteArray() : null;
    if (uak == null) {
      logger.error("directoryInsertEntryFromItem uak null,pni:" + UUID.nameUUIDFromBytes(pni));
    } else {
      logger.info("directoryInsertEntryFromItem e164:" + e164);
    }
    return new DirectoryEntry(e164, aci, pni, uak);
//    } else {
//      return null;
//    }
  }

  @VisibleForTesting
  static long e164FromString(final String s) {
    if (!s.startsWith("+")) {
      throw new IllegalArgumentException("e164 not prefixed with '+'");
    }

    return Long.parseLong(s, 1, s.length(), 10);
  }

  @Override
  public boolean hasFinishedInitialAccountPopulation() {
    return accountsPopulated.get();
  }

  @Override
  public boolean isHealthy() {
    return true;
  }
}
